package org.lc.gateway.util;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.ObjUtil;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import jakarta.servlet.http.HttpServletRequest;
import org.lc.platform.base.constant.AuthConstant;
import org.lc.platform.base.domain.JwtUser;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.oauth2.jwt.*;
import reactor.core.publisher.Mono;

import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

public class JwtTokenUtil {

    public static ReactiveJwtDecoder getJwtDecoder(String jwkSetString) {
        return new NimbusReactiveJwtDecoder(Objects.requireNonNull(getRSAPublicKey(jwkSetString)));

    }

    public static String getJwtFromRequest(ServerHttpRequest request) {

        String bearerToken = request.getHeaders().getFirst(AuthConstant.JWT_TOKEN_HEADER);
        if (bearerToken == null) {
            return null;
        }
        if (bearerToken.startsWith(AuthConstant.JWT_TOKEN_PREFIX_BEARER)) {
            return bearerToken.substring(7);
        }
        return null;
    }

    public static JwtUser getJwtUser(ReactiveJwtDecoder jwtDecoder, String token) {
        var userMap = getJwtMapUser(jwtDecoder, token);
        JwtUser user = new JwtUser();
        BeanUtil.copyProperties(userMap, user);
        return user;
    }

    public static Map<String, Object> getJwtMapUser(ReactiveJwtDecoder jwtDecoder, String token) {
        try {
            if (ObjUtil.isNotNull(jwtDecoder)) {
                Mono<Jwt> jwt = jwtDecoder.decode(token);
                AtomicReference<Map<String, Object>> claims = new AtomicReference<>();
                if (jwt != null) {
                    Mono<Map<String, Object>> claimsMono = jwt.flatMap(j -> Mono.just(j.getClaims()));
                    claimsMono.subscribe(claims::set);
                }
                return claims.get();
            }
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
        return new HashMap<>();
    }

    private static RSAPublicKey getRSAPublicKey(String jwkSetString) {
        try {
            JWKSet jwkSet = JWKSet.parse(jwkSetString);
            List<JWK> jwkList = jwkSet.getKeys();
            for (JWK jwk : jwkList) {
                var jwkRSA = jwk.toRSAKey();
                if (jwkRSA.isPrivate()) {
                    return jwkRSA.toRSAPublicKey();
                }
            }
        } catch (ParseException | JOSEException e) {
            System.out.println(e.getMessage());
        }
        return null;
    }

    public static String getClientSign(ServerHttpRequest request) {
        if (request == null) {
            return null;
        }
        String remoteAddr = IpUtils.getIP(request);
        String userAgent = request.getHeaders().getFirst("User-Agent");
        return Base64.decodeStr(userAgent + ";" + remoteAddr, StandardCharsets.UTF_8);
    }
}
